#include "bplus_tree.h"
#include <cstring>
#include <algorithm>

using namespace std;

BPlusTree::BPlusTree(const string& index_filename) : min_keys((MAX_KEYS + 1) / 2) {
    // 同理于数据文件存在则以读写追加的方式打开，不存在则创建再打开
    index_file.open(index_filename, ios::binary | ios::in | ios::out | ios::ate);
    if (!index_file.is_open()) {
        index_file.open(index_filename, ios::binary | ios::out);
        index_file.close();
        index_file.open(index_filename, ios::binary | ios::in | ios::out | ios::ate);
        // 初始的根节点就是叶子结点
        root = allocate_new_node();
        BPlusTreeNode root_node;
        root_node.is_leaf = true;
        //写入索引文件偏移量为root的地方
        write_node(root, root_node);
    } else {
        // 存在但为空文件
        if (index_file.tellg() == 0) {
            root = allocate_new_node();
            BPlusTreeNode root_node;
            root_node.is_leaf = true;
            write_node(root, root_node);
        } else {
            // 首地址是根节点的位置
            index_file.seekg(0, ios::beg);
            index_file.read(reinterpret_cast<char*>(&root), sizeof(root));
        }
    }
}

BPlusTree::~BPlusTree() {
    if (index_file.is_open()) {
        // 关闭B+树但是需要记录其根节点的偏移量，以便下次用此索引文件来初始化B+树时能直接找到根节点
        // 找到根节点就能还原出整棵树
        index_file.seekp(0, ios::beg);
        index_file.write(reinterpret_cast<const char*>(&root), sizeof(root));
        index_file.close();
    }
}

// 写到索引文件中偏移量为offset的位置
void BPlusTree::write_node(long long offset, const BPlusTreeNode& node) {
    index_file.seekp(offset, ios::beg);
    index_file.write(reinterpret_cast<const char*>(&node.is_leaf), sizeof(node.is_leaf));
    index_file.write(reinterpret_cast<const char*>(&node.key_count), sizeof(node.key_count));
    index_file.write(reinterpret_cast<const char*>(node.keys), sizeof(int) * MAX_KEYS);
    index_file.write(reinterpret_cast<const char*>(node.children), sizeof(long long) * (MAX_KEYS + 1));
    if (node.is_leaf) {
        index_file.write(reinterpret_cast<const char*>(&node.next_leaf), sizeof(node.next_leaf));
    }
}

// 读索引文件中偏移量为offset的位置
void BPlusTree::read_node(long long offset, BPlusTreeNode& node) {
    index_file.seekg(offset, ios::beg);
    index_file.read(reinterpret_cast<char*>(&node.is_leaf), sizeof(node.is_leaf));
    index_file.read(reinterpret_cast<char*>(&node.key_count), sizeof(node.key_count));
    index_file.read(reinterpret_cast<char*>(node.keys), sizeof(int) * MAX_KEYS);
    index_file.read(reinterpret_cast<char*>(node.children), sizeof(long long) * (MAX_KEYS + 1));
    if (node.is_leaf) {
        index_file.read(reinterpret_cast<char*>(&node.next_leaf), sizeof(node.next_leaf));
    }
}

// 分配一个新节点并返回此新结点在索引文件中的偏移量
long long BPlusTree::allocate_new_node() {
    // 调整写入指针至结尾，即新结点一定插入在索引文件的结尾
    index_file.seekp(0, ios::end);
    long long offset = index_file.tellp();
    BPlusTreeNode empty_node;
    // 将新结点写入至索引文件的offset位置
    write_node(offset, empty_node);
    return offset;
}

// 插入(key,data_offset)到某个叶子结点中
void BPlusTree::insert_into_leaf(BPlusTreeNode& leaf, int key, long long data_offset) {
    int i = leaf.key_count - 1;
    // 模拟插入排序
    while (i >= 0 && leaf.keys[i] > key) {
        leaf.keys[i + 1] = leaf.keys[i];
        leaf.children[i + 1] = leaf.children[i];
        i--;
    }
    leaf.keys[i + 1] = key;
    leaf.children[i + 1] = data_offset;
    leaf.key_count++;
}

// 一个结点满了则开始分裂
// 将一半的数据给兄弟结点
void BPlusTree::split_leaf(BPlusTreeNode& leaf, long long leaf_offset) {
    long long new_leaf_offset = allocate_new_node();
    BPlusTreeNode new_leaf;
    new_leaf.is_leaf = true;
    // 形似链表，只不过这里不是指针，而是用偏移量来连接结点
    new_leaf.next_leaf = leaf.next_leaf;
    leaf.next_leaf = new_leaf_offset;

    // 分一半数据给兄弟，同时孩子也要分一半给兄弟
    int split_pos = leaf.key_count / 2;
    new_leaf.key_count = leaf.key_count - split_pos;
    for (int i = 0; i < new_leaf.key_count; i++) {
        new_leaf.keys[i] = leaf.keys[split_pos + i];
        new_leaf.children[i] = leaf.children[split_pos + i];
    }
    leaf.key_count = split_pos;

    write_node(leaf_offset, leaf);
    write_node(new_leaf_offset, new_leaf);
}

// 插入到非叶子结点，插入规则也是类似插入排序，这里还需要更新孩子
void BPlusTree::insert_into_internal(BPlusTreeNode& parent, int key, long long left, long long right) {
    int i = parent.key_count - 1;
    while (i >= 0 && parent.keys[i] > key) {
        parent.keys[i + 1] = parent.keys[i];
        parent.children[i + 2] = parent.children[i + 1];
        i--;
    }
    parent.keys[i + 1] = key;
    //更新左右孩子
    parent.children[i + 1] = left;
    parent.children[i + 2] = right;
    parent.key_count++;
}

// 非叶子结点的分裂，这里需要注意如果是根节点在分裂则需要生成一个新的根节点
void BPlusTree::split_internal(BPlusTreeNode& internal, long long internal_offset) {
    long long new_internal_offset = allocate_new_node();
    BPlusTreeNode new_internal;
    new_internal.is_leaf = false;

    int split_pos = internal.key_count / 2;
    int promoted_key = internal.keys[split_pos];

    new_internal.key_count = internal.key_count - split_pos - 1;
    for (int i = 0; i < new_internal.key_count; i++) {
        new_internal.keys[i] = internal.keys[split_pos + 1 + i];
        new_internal.children[i] = internal.children[split_pos + 1 + i];
    }
    new_internal.children[new_internal.key_count] = internal.children[internal.key_count];
    internal.key_count = split_pos;

    write_node(internal_offset, internal);
    write_node(new_internal_offset, new_internal);

    // 生成新的根节点并初始化，重要的是更新根节点的偏移量root
    if (internal_offset == root) {
        long long new_root_offset = allocate_new_node();
        BPlusTreeNode new_root;
        new_root.is_leaf = false;
        new_root.key_count = 1;
        new_root.keys[0] = promoted_key;
        new_root.children[0] = internal_offset;
        new_root.children[1] = new_internal_offset;
        root = new_root_offset;
        write_node(new_root_offset, new_root);
    }
}

//递归从根节点开始向叶子结点插入(key,data_offset)键值对
//(promoted_key,promoted_node)为下层分裂后需要传给上层的信息，即插入到上层的新结点
bool BPlusTree::insert_recursive(long long node_offset, int key, long long data_offset, int& promoted_key, long long& promoted_node) {
    BPlusTreeNode node;
    read_node(node_offset, node);

    // 叶子结点为递归的最底层，可以视作递归出口
    if (node.is_leaf) {
        insert_into_leaf(node, key, data_offset);
        // 没有超过键数量的最大值则不分裂
        if (node.key_count < MAX_KEYS) {
            write_node(node_offset, node);
            return false;
        } else {
            split_leaf(node, node_offset);
            // 此时的key_count就是兄弟节点的第一个值，需要传给父亲结点
            promoted_key = node.keys[node.key_count];
            // 此时的next_leaf指向的就是兄弟结点的偏移量（地址）
            promoted_node = node.next_leaf;
            read_node(node_offset, node);
            write_node(node_offset, node);
            return true;
        }
    } else {
        int i = 0;
        // 找到下一层孩子的位置
        while (i < node.key_count && node.keys[i] < key) i++;
        long long child_offset = node.children[i];

        int child_promoted_key;
        long long child_promoted_node;
        // 拿到该去的孩子的信息向下层递归
        bool child_split = insert_recursive(child_offset, key, data_offset, child_promoted_key, child_promoted_node);

        //如果孩子没有分类返回false，提醒上层并未分裂
        if (!child_split) return false;

        // 如果分裂了，将子节点给出的需要提升到父亲的新数据插入到非叶子结点
        // 左孩子是原来的孩子，右孩子为分裂出来的兄弟
        insert_into_internal(node, child_promoted_key, child_offset, child_promoted_node);
        //非叶子结点也一样，没有超过可以容纳键的最大值则直接返回false告诉上层本层没有分裂
        if (node.key_count < MAX_KEYS) {
            write_node(node_offset, node);
            return false;
        } else {
            //分裂后将新结点的信息返回给上层
            split_internal(node, node_offset);
            promoted_key = node.keys[node.key_count];
            promoted_node = node.children[node.key_count + 1];
            read_node(node_offset, node);
            write_node(node_offset, node);
            return true;
        }
    }
}

//真正暴露给外面插入B+树的接口，刚刚所有其他插入接口为此服务
void BPlusTree::insert(int key, long long data_offset) {
    int promoted_key;
    long long promoted_node;
    //仍然查看原根节点是否分裂
    bool root_split = insert_recursive(root, key, data_offset, promoted_key, promoted_node);
    if (root_split) {
        //若分裂则创造新的根节点，并连接原根节点和分裂出的兄弟结点，重要的是更新root偏移量，以便于找到root
        long long new_root_offset = allocate_new_node();
        BPlusTreeNode new_root;
        new_root.is_leaf = false;
        new_root.key_count = 1;
        new_root.keys[0] = promoted_key;
        new_root.children[0] = root;
        new_root.children[1] = promoted_node;
        root = new_root_offset;
        write_node(new_root_offset, new_root);
    }
}

// 实现内部节点重分配
// 主要用在删除后若某一个叶子结点的结点数目不足最小值时，去找它左边的兄弟节点借而不是直接删除
void BPlusTree::redistribute_internal(BPlusTreeNode& left, BPlusTreeNode& right, long long left_offset, long long right_offset) {
    //腾出位置给从左兄弟借来的结点
    right.children[right.key_count + 1] = right.children[right.key_count];
    for (int i = right.key_count; i > 0; i--) {
        right.keys[i] = right.keys[i - 1];
        right.children[i] = right.children[i - 1];
    }
    right.keys[0] = left.keys[left.key_count - 1];
    right.children[0] = left.children[left.key_count];
    right.key_count++;
    left.key_count--;
    write_node(left_offset, left);
    write_node(right_offset, right);
}

// 当无法再分配时，即左兄弟的键的数量已经等于了最小值，这个时候只能合并
void BPlusTree::merge_internal(BPlusTreeNode& left, BPlusTreeNode& right, long long left_offset, long long right_offset, int key) {
    //找到在父亲结点中区分两者的key
    left.keys[left.key_count] = key;
    //开始合并
    for (int i = 0; i < right.key_count; i++) {
        left.keys[left.key_count + 1 + i] = right.keys[i];
        left.children[left.key_count + 1 + i] = right.children[i];
    }
    left.children[left.key_count + 1 + right.key_count] = right.children[right.key_count];
    left.key_count += right.key_count + 1;
    write_node(left_offset, left);
}

void BPlusTree::redistribute_leaf(BPlusTreeNode& left, BPlusTreeNode& right, long long left_offset, long long right_offset) {
    right.key_count++;
    for (int i = right.key_count - 1; i > 0; i--) {
        right.keys[i] = right.keys[i - 1];
        right.children[i] = right.children[i - 1];
    }
    right.keys[0] = left.keys[left.key_count - 1];
    right.children[0] = left.children[left.key_count - 1];
    left.key_count--;
    write_node(left_offset, left);
    write_node(right_offset, right);
}

void BPlusTree::merge_leaf(BPlusTreeNode& left, BPlusTreeNode& right, long long left_offset, long long right_offset) {
    for (int i = 0; i < right.key_count; i++) {
        left.keys[left.key_count + i] = right.keys[i];
        left.children[left.key_count + i] = right.children[i];
    }
    left.key_count += right.key_count;
    left.next_leaf = right.next_leaf;
    write_node(left_offset, left);
}

//在某个叶子上删除一个值
bool BPlusTree::remove_from_leaf(BPlusTreeNode& leaf, long long leaf_offset, int key) {
    int i = 0;
    while (i < leaf.key_count && leaf.keys[i] != key) i++;
    if (i >= leaf.key_count) return false;

    for (int j = i; j < leaf.key_count - 1; j++) {
        leaf.keys[j] = leaf.keys[j + 1];
        leaf.children[j] = leaf.children[j + 1];
    }
    leaf.key_count--;
    write_node(leaf_offset, leaf);
    //为真说明需要借或者合并
    return leaf.key_count < min_keys;
}

// 从根节点开始搜索递归删除
bool BPlusTree::remove_recursive(long long node_offset, int key) {
    BPlusTreeNode node;
    read_node(node_offset, node);

    //递归出口如果时叶子结点直接执行删除
    if (node.is_leaf) {
        return remove_from_leaf(node, node_offset, key);
    }

    int i = 0;
    while (i < node.key_count && node.keys[i] < key) i++;
    long long child_offset = node.children[i];
    bool child_underflow = remove_recursive(child_offset, key);

    if (!child_underflow) return false;

    BPlusTreeNode child_node;
    read_node(child_offset, child_node);
    if (child_node.key_count >= min_keys) return false;

    BPlusTreeNode sibling_node;
    long long sibling_offset = -1;
    int sibling_pos = -1;

    //判断是否有左兄弟，有左兄弟就找左兄弟，没有就找右兄弟
    if (i > 0) {
        sibling_offset = node.children[i - 1];
        sibling_pos = i - 1;
    } else if (i < node.key_count) {
        sibling_offset = node.children[i + 1];
        sibling_pos = i + 1;
    }

    read_node(sibling_offset, sibling_node);
    if (sibling_node.key_count > min_keys) {
        // 兄弟有多余结点是最好，只需要借就可以
        if (sibling_pos == i - 1) {
            redistribute_internal(sibling_node, child_node, sibling_offset, child_offset);
            node.keys[i - 1] = sibling_node.keys[sibling_node.key_count - 1];
        } else {
            redistribute_internal(child_node, sibling_node, child_offset, sibling_offset);
            node.keys[i] = sibling_node.keys[0];
        }
        write_node(node_offset, node);
        return false;
    } else {
        // 若兄弟的结点也不够就需要合并
        if (sibling_pos == i - 1) {
            //与左兄弟合并
            merge_internal(sibling_node, child_node, sibling_offset, child_offset, node.keys[i - 1]);
            for (int j = i - 1; j < node.key_count - 1; j++) {
                node.keys[j] = node.keys[j + 1];
                node.children[j + 1] = node.children[j + 2];
            }
            node.key_count--;
        } else {
            //与右兄弟合并
            merge_internal(child_node, sibling_node, child_offset, sibling_offset, node.keys[i]);
            for (int j = i; j < node.key_count - 1; j++) {
                node.keys[j] = node.keys[j + 1];
                node.children[j + 1] = node.children[j + 2];
            }
            node.key_count--;
        }
        write_node(node_offset, node);
        // 判断合并因为删除了分割键看是否仍然需要做出借或者合并的调整
        return node.key_count < min_keys;
    }
}

// 真正向外展示的删除的接口
bool BPlusTree::remove(int key) {
    bool root_underflow = remove_recursive(root, key);
    BPlusTreeNode root_node;
    read_node(root, root_node);

    if (root_underflow && !root_node.is_leaf) {
        root = root_node.children[0];
    }
    return true;
}

bool BPlusTree::find(int key, long long& data_offset) {
    //从根节点开始向下搜索
    long long current_offset = root;
    BPlusTreeNode current_node;

    while (true) {
        read_node(current_offset, current_node);
        if (current_node.is_leaf) {
            // 到叶子结点说明已经找到正确范围，燃后进去搜索，如果搜索到了说明存在，否则就是不存在
            for (int i = 0; i < current_node.key_count; i++) {
                if (current_node.keys[i] == key) {
                    data_offset = current_node.children[i];
                    return true;
                }
            }
            return false;
        }
        int i = 0;
        //非叶子结点搜索，找到在哪个孩子那里
        while (i < current_node.key_count && current_node.keys[i] < key) i++;
        current_offset = current_node.children[i];
    }
}

vector<long long> BPlusTree::range_query(int start, int end) {
    vector<long long> result;
    long long current_offset = root;
    BPlusTreeNode current_node;

    //找到start这个key对应的位置，从此处开始向后寻找直到找到>=end说明所有结果都已经被找到
    while (true) {
        read_node(current_offset, current_node);
        if (current_node.is_leaf) break;
        int i = 0;
        while (i < current_node.key_count && current_node.keys[i] < start) i++;
        current_offset = current_node.children[i];
    }

    //搜寻叶子结点直至找到 >=end
    while (current_offset != -1) {
        read_node(current_offset, current_node);
        for (int i = 0; i < current_node.key_count; i++) {
            if (current_node.keys[i] >= end) return result;
            if (current_node.keys[i] >= start) {
                result.push_back(current_node.children[i]);
            }
        }
        //像链表一般，这里保存的是下一个结点的偏移量
        current_offset = current_node.next_leaf;
    }
    return result;
}